#include "RgDetourScorer.hh"

#include "G4Alpha.hh"
#include "G4Step.hh"
#include "G4Track.hh"
#include "G4UnitsTable.hh"
//#include "G4SystemOfUnits.hh"
// #include "G4ParticleDefinition.hh"

RgDetourScorer::RgDetourScorer(G4String name, G4int depth, G4int weight) : G4VPrimitiveScorer(name, depth), fDepth(depth), fWeight(weight)
{
    _fAlpha = G4Alpha::Alpha();
    // 声明我们将会产生 3 个 map
    // 名字分别叫  <primitiveName>/CSDA  <primitiveName>/Proj  <primitiveName>/Detour
    // 这些名字在 macro 里会被用到
}

G4bool RgDetourScorer::ProcessHits(G4Step* step, G4TouchableHistory*)
{
    auto* const track = step->GetTrack();
    auto* const partDef = track->GetParticleDefinition();

    // ---- 1. 过滤：只记录主粒子，且只记录 α（可改成任何粒子） ----
    if (track->GetParentID() != 0) return TRUE;  // 只盯 primary
    // if (partDef != _fAlpha) return true;

    const G4int tkID = track->GetTrackID();
    // 第一次见这条 track 就新开一个 slot
    if (tkID >= (G4int)_fPerTrack.size()) _fPerTrack.resize(tkID + 1);

    TrackData& d = _fPerTrack[tkID];

    // ---- 2. 第一次 step 记下起点 ----
    if (track->GetCurrentStepNumber() == 1)
    {
        d.start = track->GetPosition();
        d.csda = 0.0;
    }

    // ---- 3. 每一步累加 CSDA ----
    d.csda += step->GetStepLength();

    // ---- 4. 粒子停掉或离开世界体积 → 写结果 ----
    G4bool stopped = (track->GetTrackStatus() == fStopAndKill) || (track->GetNextVolume() == nullptr);

    if (stopped)
    {
        d.end = track->GetPosition();

        // 投影轴取 Z（可改成 User 指定的 axis）
        const G4double proj = std::abs(d.end.z() - d.start.z());

        // 防止 0 除
        const G4double detour = (d.csda > 0.0) ? proj / d.csda : 1.0;

        // 5. 填 hits map（index=0 代表“整个 mesh 的唯一 bin”）
        G4int idx = 0;
        _fCSDAMap->add(idx, d.csda);
        _fProjMap->add(idx, proj);
        _fDetourMap->add(idx, detour);
    }
    return true;
}

// ----------- 以下 boilerplate 直接抄即可 -----------
void RgDetourScorer::Initialize(G4HCofThisEvent* HCE)
{
    //auto* const dctt = this->GetMultiFunctionalDetector();
    const auto& dcttName = this->GetMultiFunctionalDetector()->GetName();
    const auto& primName = this->GetName();
    // detector name, collection name

    // _fCSDAMap = new G4THitsMap<G4double>(dcttName, primName);
    _fCSDAMap = new G4THitsMap<G4double>(dcttName, primName + "/CSDA");
    _fProjMap = new G4THitsMap<G4double>(dcttName, primName + "/Proj");
    _fDetourMap = new G4THitsMap<G4double>(dcttName, primName);

    if (_fHCID[2] < 0)
    {
        _fHCID[2] = GetCollectionID(0);
    }
    HCE->AddHitsCollection(_fHCID[2], _fDetourMap);
}

void RgDetourScorer::EndOfEvent(G4HCofThisEvent*)
{
    /* 合并已在 add() 时完成，无需额外动作 */

#if 0
    auto lPrint = [](G4THitsMap<G4double>* m, G4String title)
    {
        G4cout << "-------- " << title << " --------" << G4endl;
        std::map<G4int, G4double*>::iterator it;
        for (it = m->GetMap()->begin(); it != m->GetMap()->end(); ++it)
            G4cout << " idx=" << it->first << "  val=" << G4BestUnit(*(it->second), "Length") << G4endl;
    };
    lPrint(_fCSDAMap, "CSDA Range");
#endif
}

void RgDetourScorer::PrintAll()
{
    auto print = [](G4THitsMap<G4double>* m, G4String title)
    {
        G4cout << "-------- " << title << " --------" << G4endl;
        std::map<G4int, G4double*>::iterator it;
        for (it = m->GetMap()->begin(); it != m->GetMap()->end(); ++it)
            G4cout << " idx=" << it->first << "  val=" << G4BestUnit(*(it->second), "Length") << G4endl;
    };
    print(_fCSDAMap, "CSDA Range");
    print(_fProjMap, "Projected Range");
    print(_fDetourMap, "Detour Factor");
}